#include "inflation_layer.h"


namespace costmap_2d{

InflationLayer::InflationLayer(): resolution_(0)
    , inflation_radius_(0)
    , inscribed_radius_(0)
    , weight_(0)
    , inflate_unknown_(false)
    , cell_inflation_radius_(0)
    , cached_cell_inflation_radius_(0)
    , seen_(NULL)
    , cached_costs_(NULL)
    , cached_distances_(NULL)
    , last_min_x_(-std::numeric_limits<float>::max())
    , last_min_y_(-std::numeric_limits<float>::max())
    , last_max_x_(std::numeric_limits<float>::max())
    , last_max_y_(std::numeric_limits<float>::max()){
    
    need_reinflation_ = true;
}

InflationLayer::~InflationLayer(){
    deleteKernels();
    if (seen_)
        delete[] seen_;
}

void InflationLayer::onInitialize(){
    if (seen_)
        delete[] seen_;
    seen_ = NULL;
    seen_size_ = 0;  
    matchSize(); 
}

void InflationLayer::matchSize(){
    setInflationParameters(0.2, 10.0);
    costmap_2d::Costmap2D* costmap = layered_costmap_->getCostmap2D();
    resolution_ = costmap->getResolution();
    cell_inflation_radius_ = cellDistance(inflation_radius_);
    computeCaches();

    unsigned int size_x = costmap->getWidth(), size_y = costmap->getHeight();
    if (seen_)
        delete[] seen_;
    seen_size_ = size_x * size_y;
    seen_ = new bool[seen_size_];
}

unsigned char InflationLayer::computeCost(double distance) const{
    unsigned char cost = 0;
    if (distance == 0)
        cost = LETHAL_OBSTACLE;
    else if (distance * resolution_ <= inscribed_radius_)
        cost = INSCRIBED_INFLATED_OBSTACLE;
    else
    {
        // make sure cost falls off by Euclidean distance
        double euclidean_distance = distance * resolution_;
        double factor = exp(-1.0 * weight_ * (euclidean_distance - inscribed_radius_));
        cost = (unsigned char)((INSCRIBED_INFLATED_OBSTACLE - 1) * factor);
    }
    return cost;
}

double InflationLayer::distanceLookup(int mx, int my, int src_x, int src_y){
    unsigned int dx = abs(mx - src_x);
    unsigned int dy = abs(my - src_y);
    return cached_distances_[dx][dy];
}

unsigned char InflationLayer::costLookup(int mx, int my, int src_x, int src_y){
    unsigned int dx = abs(mx - src_x);
    unsigned int dy = abs(my - src_y);
    return cached_costs_[dx][dy];
}

void InflationLayer::updateBounds(double robot_x, double robot_y, double robot_yaw, double* min_x,
    double* min_y, double* max_x, double* max_y){
    if (need_reinflation_)
    {
        last_min_x_ = *min_x;
        last_min_y_ = *min_y;
        last_max_x_ = *max_x;
        last_max_y_ = *max_y;
        // For some reason when I make these -<double>::max() it does not
        // work with Costmap2D::worldToMapEnforceBounds(), so I'm using
        // -<float>::max() instead.
        *min_x = -std::numeric_limits<float>::max();
        *min_y = -std::numeric_limits<float>::max();
        *max_x = std::numeric_limits<float>::max();
        *max_y = std::numeric_limits<float>::max();
        need_reinflation_ = false;
    }
    else
    {
        double tmp_min_x = last_min_x_;
        double tmp_min_y = last_min_y_;
        double tmp_max_x = last_max_x_;
        double tmp_max_y = last_max_y_;
        last_min_x_ = *min_x;
        last_min_y_ = *min_y;
        last_max_x_ = *max_x;
        last_max_y_ = *max_y;
        *min_x = std::min(tmp_min_x, *min_x) - inflation_radius_;
        *min_y = std::min(tmp_min_y, *min_y) - inflation_radius_;
        *max_x = std::max(tmp_max_x, *max_x) + inflation_radius_;
        *max_y = std::max(tmp_max_y, *max_y) + inflation_radius_;
    }
}

void InflationLayer::updateCosts(costmap_2d::Costmap2D& master_grid, int min_i, int min_j, int max_i, int max_j)
{
//   boost::unique_lock < boost::recursive_mutex > lock(*inflation_access_);
  if (cell_inflation_radius_ == 0)
    return;

  // make sure the inflation list is empty at the beginning of the cycle (should always be true)
  ROS_ASSERT_MSG(inflation_cells_.empty(), "The inflation list must be empty at the beginning of inflation");

  unsigned char* master_array = master_grid.getCharMap();
  unsigned int size_x = master_grid.getWidth(), size_y = master_grid.getHeight();

  if (seen_ == NULL) {
    ROS_WARN("InflationLayer::updateCosts(): seen_ array is NULL");
    seen_size_ = size_x * size_y;
    seen_ = new bool[seen_size_];
  }
  else if (seen_size_ != size_x * size_y)
  {
    ROS_WARN("InflationLayer::updateCosts(): seen_ array size is wrong");
    delete[] seen_;
    seen_size_ = size_x * size_y;
    seen_ = new bool[seen_size_];
  }
  memset(seen_, false, size_x * size_y * sizeof(bool));

  // We need to include in the inflation cells outside the bounding
  // box min_i...max_j, by the amount cell_inflation_radius_.  Cells
  // up to that distance outside the box can still influence the costs
  // stored in cells inside the box.
  min_i -= cell_inflation_radius_;
  min_j -= cell_inflation_radius_;
  max_i += cell_inflation_radius_;
  max_j += cell_inflation_radius_;

  min_i = std::max(0, min_i);
  min_j = std::max(0, min_j);
  max_i = std::min(int(size_x), max_i);
  max_j = std::min(int(size_y), max_j);

  // Inflation list; we append cells to visit in a list associated with its distance to the nearest obstacle
  // We use a map<distance, list> to emulate the priority queue used before, with a notable performance boost

  // Start with lethal obstacles: by definition distance is 0.0
  std::vector<CellData>& obs_bin = inflation_cells_[0.0];
  for (int j = min_j; j < max_j; j++)
  {
    for (int i = min_i; i < max_i; i++)
    {
      int index = master_grid.getIndex(i, j);
      unsigned char cost = master_array[index];
      if (cost == LETHAL_OBSTACLE)
      {
        obs_bin.push_back(CellData(index, i, j, i, j));
      }
    }
  }

  // Process cells by increasing distance; new cells are appended to the corresponding distance bin, so they
  // can overtake previously inserted but farther away cells
  std::map<double, std::vector<CellData> >::iterator bin;
  for (bin = inflation_cells_.begin(); bin != inflation_cells_.end(); ++bin)
  {
    for (int i = 0; i < bin->second.size(); ++i)
    {
      // process all cells at distance dist_bin.first
      const CellData& cell = bin->second[i];

      unsigned int index = cell.index_;

      // ignore if already visited
      if (seen_[index])
      {
        continue;
      }

      seen_[index] = true;

      unsigned int mx = cell.x_;
      unsigned int my = cell.y_;
      unsigned int sx = cell.src_x_;
      unsigned int sy = cell.src_y_;

      // assign the cost associated with the distance from an obstacle to the cell
      unsigned char cost = costLookup(mx, my, sx, sy);
      unsigned char old_cost = master_array[index];
      if (old_cost == NO_INFORMATION && (inflate_unknown_ ? (cost > FREE_SPACE) : (cost >= INSCRIBED_INFLATED_OBSTACLE)))
        master_array[index] = cost;
      else
        master_array[index] = std::max(old_cost, cost);

      // attempt to put the neighbors of the current cell onto the inflation list
      if (mx > 0)
        enqueue(index - 1, mx - 1, my, sx, sy);
      if (my > 0)
        enqueue(index - size_x, mx, my - 1, sx, sy);
      if (mx < size_x - 1)
        enqueue(index + 1, mx + 1, my, sx, sy);
      if (my < size_y - 1)
        enqueue(index + size_x, mx, my + 1, sx, sy);
    }
  }

  inflation_cells_.clear();
}


inline void InflationLayer::enqueue(unsigned int index, unsigned int mx, unsigned int my,
                                    unsigned int src_x, unsigned int src_y)
{
  if (!seen_[index])
  {
    // we compute our distance table one cell further than the inflation radius dictates so we can make the check below
    double distance = distanceLookup(mx, my, src_x, src_y);

    // we only want to put the cell in the list if it is within the inflation radius of the obstacle point
    if (distance > cell_inflation_radius_)
      return;

    // push the cell data onto the inflation list and mark
    inflation_cells_[distance].push_back(CellData(index, mx, my, src_x, src_y));
  }
}

void InflationLayer::computeCaches()
{
    LOG(INFO) << "computeCaches...... start"; 
    if (cell_inflation_radius_ == 0)
        return;

    // based on the inflation radius... compute distance and cost caches
    if (cell_inflation_radius_ != cached_cell_inflation_radius_)
    {
        deleteKernels();

        cached_costs_ = new unsigned char*[cell_inflation_radius_ + 2];
        cached_distances_ = new double*[cell_inflation_radius_ + 2];

        for (unsigned int i = 0; i <= cell_inflation_radius_ + 1; ++i)
        {
        cached_costs_[i] = new unsigned char[cell_inflation_radius_ + 2];
        cached_distances_[i] = new double[cell_inflation_radius_ + 2];
        for (unsigned int j = 0; j <= cell_inflation_radius_ + 1; ++j)
        {
            cached_distances_[i][j] = hypot(i, j);
        }
        }

        cached_cell_inflation_radius_ = cell_inflation_radius_;
    }

    for (unsigned int i = 0; i <= cell_inflation_radius_ + 1; ++i)
    {
        for (unsigned int j = 0; j <= cell_inflation_radius_ + 1; ++j)
        {
        cached_costs_[i][j] = computeCost(cached_distances_[i][j]);
        }
    }
}

void InflationLayer::deleteKernels()
{
  if (cached_distances_ != NULL)
  {
    for (unsigned int i = 0; i <= cached_cell_inflation_radius_ + 1; ++i)
    {
      if (cached_distances_[i])
        delete[] cached_distances_[i];
    }
    if (cached_distances_)
      delete[] cached_distances_;
    cached_distances_ = NULL;
  }

  if (cached_costs_ != NULL)
  {
    for (unsigned int i = 0; i <= cached_cell_inflation_radius_ + 1; ++i)
    {
      if (cached_costs_[i])
        delete[] cached_costs_[i];
    }
    delete[] cached_costs_;
    cached_costs_ = NULL;
  }
}

void InflationLayer::setInflationParameters(double inflation_radius, double cost_scaling_factor){
    LOG(INFO) << "setInflationParameters inflation_radius: " << inflation_radius 
        << " cost_scaling_factor: " << cost_scaling_factor;
    inflation_radius_ = inflation_radius;
    cell_inflation_radius_ = cellDistance(inflation_radius_);
    weight_ = cost_scaling_factor;
    need_reinflation_ = true;
    computeCaches();
    LOG(INFO) << "setInflationParameters...... over";
}

}  // namespace costmap_2d
